#include <slam_2d_ros/icp_2d/likelihood_filed.h>

namespace slam_2d{


LikelihoodField::LikelihoodField(){

    // 构建模板
    const int range = 20;  // 生成多少个像素的模板
    for (int x = -range; x <= range; ++x) {
        for (int y = -range; y <= range; ++y) {
            model_.emplace_back(x, y, std::sqrt((x * x) + (y * y)));
        }
    }
}


// 设置被配准的那个scan
void LikelihoodField::SetSourceScan(Scan2DPtr scan){
    source_ = scan;
}

// 设置pose 
void LikelihoodField::SetPose(const Pose2D& pose){
    pose_ = pose;
}

// 从占据栅格地图生成一个似然场地图
void LikelihoodField::SetFieldImageFromOccuMap(const cv::Mat& occu_map){
    const int boarder = 25;
    field_ = cv::Mat(1000, 1000, CV_32F, 30.0);

    for (int x = boarder; x < occu_map.cols - boarder; ++x) {
        for (int y = boarder; y < occu_map.rows - boarder; ++y) {
            if (occu_map.at<uchar>(y, x) < 127) {
                // 在该点生成一个model
                for (auto& model_pt : model_) {
                    int xx = int(x + model_pt.dx_);
                    int yy = int(y + model_pt.dy_);
                    if (xx >= 0 && xx < field_.cols && yy >= 0 && yy < field_.rows &&
                        field_.at<float>(yy, xx) > model_pt.residual_) {
                        field_.at<float>(yy, xx) = model_pt.residual_;
                        
                        // LOG(INFO) << "[SetFieldImageFromOccuMap] x " << x << ", y " << y
                        //     << ", xx " << xx << ", yy " << yy << ", residual " << model_pt.residual_;
                    }
                }
            }
            // if(fabs(x - 494) < 3 && fabs(y - 489) < 3)
            //     LOG(INFO) << "[SetFieldImageFromOccuMap] x " << x << ", y " << y
            //         << ", field value " << field_.at<float>(y, x)
            //         << ", occu_map value " << int(occu_map.at<uchar>(y, x));
        }
    }  
    // cv::imwrite("/home/xxty/trobot/trobot_ros/trobot_occu_map.png", occu_map);  
    // cv::imwrite("/home/xxty/trobot/trobot_ros/trobot_ofield.png", field_);
}

// 使用高斯牛顿法配准
bool LikelihoodField::AlignGaussNewton(Pose2D& init_pose){
    int iterations = 10;
    double cost = 0, lastCost = 0;
    Pose2D current_pose = init_pose;
    const int min_effect_pts = 20;  // 最小有效点数
    const int image_boarder = 20;   // 预留图像边界
    LOG(INFO) << "[AlignGaussNewton] current_pose " << PrintData(current_pose);
    float theta1 = Eigen::Rotation2Dd(current_pose.rotation()).angle();
    LOG(INFO) << "[AlignGaussNewton] theta = " << theta1;

    has_outside_pts_ = false;
    for (int iter = 0; iter < iterations; ++iter) {
        Mat3d H = Mat3d::Zero();
        Vec3d b = Vec3d::Zero();
        cost = 0;

        int effective_num = 0;  // 有效点数

        // 遍历 source
        for (size_t i = 0; i < source_->ranges.size(); ++i) {
            float r = source_->ranges[i];
            if (r < source_->range_min || r > source_->range_max) {
                continue;
            }

            float angle = source_->angle_min + i * source_->angle_increment;
            if (angle < source_->angle_min + 30 * M_PI / 180.0 || angle > source_->angle_max - 30 * M_PI / 180.0) {
                continue;
            }

            // 使用 Pose2D 提取旋转角
            float theta = Eigen::Rotation2Dd(current_pose.rotation()).angle();
            Vec2d pw = current_pose * Vec2d(r * std::cos(angle), r * std::sin(angle));
            // if(i % 50 == 0){
            //     LOG(INFO) << "[AlignGaussNewton] i = " << i << ", Vec2d: " << pw;
            // }

            // 在 field 中的图像坐标
            Vec2i pf = (pw * resolution_ + Vec2d(500, 500)).cast<int>();
            
            if (pf[0] >= image_boarder && pf[0] < field_.cols - image_boarder && pf[1] >= image_boarder &&
                pf[1] < field_.rows - image_boarder) {
                effective_num++;
                // LOG(INFO) << "[AlignGaussNewton] i = " << i 
                //     << ", pw: " << pw.transpose()
                //     <<" , theta: " << theta
                //     << ", angle: " << angle << ", r: " << r
                //     << ", (y, x+1): " << field_.at<float>(pf[1], pf[0] + 1)
                //     << ", (y, x-1): " << field_.at<float>(pf[1], pf[0] - 1)
                //     << ", (x, y+1): " << field_.at<float>(pf[1] + 1, pf[0])
                //     << ", (x, y-1): " << field_.at<float>(pf[1] - 1, pf[0]);
                // 图像梯度
                float dx = 0.5 * (field_.at<float>(pf[1], pf[0] + 1) - field_.at<float>(pf[1], pf[0] - 1));
                float dy = 0.5 * (field_.at<float>(pf[1] + 1, pf[0]) - field_.at<float>(pf[1] - 1, pf[0]));

                Vec3d J;
                J << resolution_ * dx, resolution_ * dy,
                    -resolution_ * dx * r * std::sin(angle + theta) + resolution_ * dy * r * std::cos(angle + theta);
                H += J * J.transpose();

                float e = field_.at<float>(pf[1], pf[0]);
                b += -J * e;

                cost += e * e;

                // LOG(INFO) << "[AlignGaussNewton] i = " << i 
                //     << ", pf: " << pf.transpose()
                //     <<" , dx: " << dx << ", dy: " << dy
                //     << ", e: " << e << ", cost: " << cost << ", lastCost: " << lastCost;
            } else {
                has_outside_pts_ = true;
            }
        }

        if (effective_num < min_effect_pts) {
            return false;
        }

        // solve for dx
        Vec3d dx = H.ldlt().solve(b);
        if (isnan(dx[0])) {
            break;
        }

        cost /= effective_num;
        if (iter > 0 && cost >= lastCost) {
            break;
        }

        // LOG(INFO) << "iter " << iter << " cost = " << cost << ", effect num: " << effective_num;

        current_pose.translation() += dx.head<2>();  // 平移更新
        Eigen::Rotation2Dd delta_rot(dx[2]);         // 构造旋转增量
        current_pose.rotate(delta_rot);              // 更新旋转部分

        lastCost = cost;
    }

    init_pose = current_pose;
    return true;
}

// 获取场函数，转换为RGB图像
cv::Mat LikelihoodField::GetFieldImage(){
    cv::Mat image(field_.rows, field_.cols, CV_8UC3);
    for (int x = 0; x < field_.cols; ++x) {
        for (int y = 0; y < field_.rows; ++y) {
            float r = field_.at<float>(y, x) * 255.0 / 30.0;
            image.at<cv::Vec3b>(y, x) = cv::Vec3b(uchar(r), uchar(r), uchar(r));
        }
    }
    return image;    
}

} // namespace slam_2d